// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/vector_builder.h>
#include <turbo/strings/string_builder.h>

namespace kumo::pollux {
    BufferPtr make_indices_in_reverse(vector_size_t size, memory::MemoryPool *pool) {
        auto indices = AlignedBuffer::allocate<vector_size_t>(size, pool);
        auto rawIndices = indices->asMutable<vector_size_t>();
        for (auto i = 0; i < size; i++) {
            rawIndices[i] = size - 1 - i;
        }
        return indices;
    }

    BufferPtr make_indices(
        vector_size_t size,
        std::function<vector_size_t(vector_size_t)> indexAt,
        memory::MemoryPool *pool) {
        BufferPtr indices = AlignedBuffer::allocate<vector_size_t>(size, pool);
        auto rawIndices = indices->asMutable<vector_size_t>();

        for (vector_size_t i = 0; i < size; i++) {
            rawIndices[i] = indexAt(i);
        }

        return indices;
    }

    VectorBuilder::~VectorBuilder() {
        // Reset the executor to wait for all the async activities to finish.
        executor_.reset();
    }

    // static
    VectorPtr VectorBuilder::wrap_in_dictionary(
        BufferPtr indices,
        vector_size_t size,
        VectorPtr vector) {
        return BaseVector::wrap_in_dictionary(
            BufferPtr(nullptr), indices, size, vector);
    }

    // static
    VectorPtr VectorBuilder::wrap_in_dictionary(
        BufferPtr indices,
        VectorPtr vector) {
        return wrap_in_dictionary(
            indices, indices->size() / sizeof(vector_size_t), vector);
    }

    BufferPtr VectorBuilder::make_odd_indices(vector_size_t size) {
        return make_indices(size, [](vector_size_t i) { return 2 * i + 1; });
    }

    BufferPtr VectorBuilder::make_even_indices(vector_size_t size) {
        return make_indices(size, [](vector_size_t i) { return 2 * i; });
    }

    BufferPtr VectorBuilder::make_indices(
        vector_size_t size,
        std::function<vector_size_t(vector_size_t)> indexAt) const {
        return ::kumo::pollux::make_indices(size, indexAt, pool());
    }

    BufferPtr VectorBuilder::make_indices(
        const std::vector<vector_size_t> &indices) const {
        auto size = indices.size();
        BufferPtr indicesBuffer =
                AlignedBuffer::allocate<vector_size_t>(size, pool());
        auto rawIndices = indicesBuffer->asMutable<vector_size_t>();

        for (int i = 0; i < size; i++) {
            rawIndices[i] = indices[i];
        }
        return indicesBuffer;
    }

    BufferPtr VectorBuilder::make_nulls(
        vector_size_t size,
        std::function<bool(vector_size_t /*row*/)> is_null_at) {
        auto nulls = AlignedBuffer::allocate<bool>(size, pool());
        auto rawNulls = nulls->asMutable<uint64_t>();
        for (auto i = 0; i < size; i++) {
            bits::set_null(rawNulls, i, is_null_at(i));
        }
        return nulls;
    }

    BufferPtr VectorBuilder::make_nulls(const std::vector<bool> &values) {
        auto nulls = allocate_nulls(values.size(), pool());
        auto rawNulls = nulls->asMutable<uint64_t>();
        for (auto i = 0; i < values.size(); i++) {
            bits::set_null(rawNulls, i, values[i]);
        }
        return nulls;
    }

    std::vector<RowVectorPtr> VectorBuilder::split(
        const RowVectorPtr &vector,
        int32_t n) {
        const auto numRows = vector->size();
        POLLUX_CHECK_GE(numRows, n);

        const auto numRowsPerVector = numRows / n;
        std::vector<RowVectorPtr> vectors;

        vector_size_t offset = 0;
        for (auto i = 0; i < n - 1; ++i) {
            vectors.push_back(std::dynamic_pointer_cast<RowVector>(
                vector->slice(offset, numRowsPerVector)));
            offset += numRowsPerVector;
        }

        vectors.push_back(std::dynamic_pointer_cast<RowVector>(
            vector->slice(offset, numRows - offset)));
        return vectors;
    }

    VectorPtr VectorBuilder::as_array(VectorPtr elements) {
        auto *pool = elements->pool();
        auto arrayType = ARRAY(elements->type());

        BufferPtr sizes =
                AlignedBuffer::allocate<vector_size_t>(1, pool, elements->size());
        BufferPtr offsets = allocateOffsets(1, pool);
        return std::make_shared<ArrayVector>(
            pool, arrayType, nullptr, 1, offsets, sizes, std::move(elements));
    }

    bool is_equal_vectors(const VectorPtr &expected, const VectorPtr &actual, std::string *err) {
        if(expected->size() != actual->size()) {
            return false;
        }
        if(!expected->type()->equivalent(*actual->type())) {
            if(err) {
                *err = turbo::str_cat("Expected ", expected->type()->toString(), ", but got ", actual->type()->toString());
            }
            return false;
        }

        for (auto i = 0; i < expected->size(); i++) {
            if(!expected->equal_value_at(actual.get(), i, i)) {

                 if(err) {
                    *err = turbo::str_cat("at ", i, ": expected ", expected->toString(i), ", but got ", actual->toString(i));
                 }
                 return false;
            }
        }
        return true;
    }

    bool is_equal_vectors(
        const VectorPtr &expected,
        const VectorPtr &actual,
        const SelectivityVector &rowsToCompare, std::string *err) {
        if(rowsToCompare.end() > actual->size()) {
             if(err) {
                    *err ="Vectors should at least have the required amount of rows that need "
                            "to be verified.";
                 }
                 return false;
        }

        if(!expected->type()->equivalent(*actual->type())) {
             if(err) {
                    *err = turbo::str_cat("Expected ", expected->type()->toString() , ", but got ", actual->type()->toString());
                 }
                 return false;
        }
        bool eq;
        rowsToCompare.applyToSelected([&](vector_size_t i) {
            if(!expected->equal_value_at(actual.get(), i, i)) {
                if(err) {
                    *err = turbo::str_cat("at " , i , ": expected " , expected->toString(i), ", but got ", actual->toString(i));
                 }
                eq = false;
            }
            eq = true;
        });
        return eq;
    }

    bool is_copyable_vector(const VectorPtr &vector) {
        try {
            auto copy =
                    BaseVector::create(vector->type(), vector->size(), vector->pool());
            copy->copy(vector.get(), 0, 0, vector->size());
            return true;
        } catch (...) {
            return false;
        }
    }
} // namespace kumo::pollux
